"""5.4 【资源模板】Resource Template 客户端开发 """
import asyncio
import json
from contextlib import AsyncExitStack
from typing import Optional

from mcp import ClientSession, ListResourcesResult
from mcp.types import ListResourceTemplatesResult
from openai import OpenAI
from mcp.client.sse import sse_client

class MCPClient:

    def __init__(self):
        #创建线程管理栈
        self.exit_stack = AsyncExitStack()
        self.session: Optional[ClientSession] = None
        self.deepseek = OpenAI(
            api_key="sk-15af4e21f828460683b16ce9e78b2346",
            base_url="https://api.deepseek.com"
        )
        self.resources = {}

    """创建连接服务端"""
    async def connect_to_server(self, server_path: str):

        # 创建sse_client
        client = sse_client(url=server_path)
        stdio_transport = await self.exit_stack.enter_async_context(client)
        read_stream, write_stream = stdio_transport

        # 三、创建ClientSession
        client_session = ClientSession(read_stream, write_stream)
        self.session = await self.exit_stack.enter_async_context(client_session)
        # 四、初始化session
        await self.session.initialize()


    async def execute(self,query:str):

        """ （一） 获取资源模板列表信息 : list_resource_templates"""
        list_resource_template_result:ListResourceTemplatesResult = await self.session.list_resource_templates()
        resources = list_resource_template_result.resourceTemplates
        print("\nConnected to server with resource:",resources)
        resource_tools = []
        for template in resources:
            # (二）修改uri改为uriTemplate
            uri = template.uriTemplate
            name = template.name
            description = template.description
            mime_type = template.mimeType
            input_schema = None
            print(f"URI: {uri}, Name: {name}, Description: {description}, Input Schema: {input_schema}")
            # 保存资源类信息，为后续function_calling做准备
            self.resources[name] = {
                "uri":uri,
                "name":name,
                "description":description,
                "mime_type":mime_type,
                "input_schema":input_schema,
            }
            # 保存资源类信息，为后续function_calling做准备
            resource_tools.append({
                "type":"function",
                "function":{
                    "name":name,
                    "description":description,
                    "parameters":input_schema
                }
            })


        # 组装function calling

        # 二、 大模型调用参数：组装messages
        messages = [
            {
                "role": "user",
                "content": query
            }
        ]
        # 第一次调用大模型做决策（选择合适的资源类）
        deepseek_response = self.deepseek.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            tools=resource_tools,
        )
        # 获取大模型的决策结果
        print("==== deepseek 决策结果：",deepseek_response)
        choice_result = deepseek_response.choices[0]
        print("最终选择的工具：",choice_result.message.tool_calls[0].function.name)

        '''  ******************* function calling 过程 *******************'''
        if choice_result.finish_reason == "tool_calls":
            # 根据大模型选择的工具，进行调用执行
            messages.append(choice_result.message.model_dump())
            # 调用工具链
            tool_call = choice_result.message.tool_calls[0]
            function_name = tool_call.function.name
            arguments = json.loads(tool_call.function.arguments)


            ''' ************** 调用资源类 *************'''
            uri = self.resources[function_name]['uri']
            # (三）修改内容：替换动态的参数
            uri = uri.format(**arguments)
            #tool_result = await self.session.call_tool(name = function_name,arguments=arguments)
            resource_result = await self.session.read_resource(uri)

            print("==== 工具调用结果元数据：",resource_result)
            # 把 resource_result.content改为resource_result.contents[0].text
            print("====  工具计算结果：",resource_result.contents[0].text)

            ''' ************** 3.5 最终推理结果 (第二次调用大模型）*************'''
            # 组装第二次调用大模型的message参数:角色需要为tool,来响应工具的调用
            messages.append({
                "role":"tool",
                "content":resource_result.contents[0].text,
                "tool_call_id":tool_call.id
            })
            #再次调用大模型
            deepseek_response = self.deepseek.chat.completions.create(
                model="deepseek-chat",
                messages=messages,
            )
            # 获取最终的结果
            result = deepseek_response.choices[0].message.content
            print("==== 最终的结果：",result)


    #关闭连接
    async def cleanup(self):
        await self.exit_stack.aclose()

async def main():
    client = MCPClient()
    try:
        await client.connect_to_server("http://localhost:8000/sse")
        await client.execute("查询book_id为123的图书信息")
    except Exception as e:
        print(f"Error: {str(e)}")
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())